package org.xhy.oap.web;

import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.action.index.IndexRequest;
import org.apache.http.HttpHost;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.xcontent.XContentType;

import java.io.IOException;

public class ElasticsearchInsertTest {

    public static void main(String[] args) throws IOException {
        try (RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(new HttpHost("localhost", 9200, "http")))) {

            // 批量请求
            BulkRequest bulkRequest = new BulkRequest();
            int batchSize = 1000; // 每次插入1000条数据
            long start = System.currentTimeMillis();
            int count = 1000005; // 初始 traceId 计数器
            for (int i = 1; i <= 10000000; i++) {
                // 构建 traceId
                String traceId = "trace_" + count++;

                // 构建 10 个 spans
                StringBuilder spansBuilder = new StringBuilder("[");
                for (int j = 1; j <= 10; j++) {
                    spansBuilder.append("{ \"spanId\": ").append(j).append(", ")
                            .append("\"traceId\": \"").append(traceId).append("\", ")
                            .append("\"parentSpanId\": ").append(j - 1).append(", ")
                            .append("\"operationName\": \"operation_").append(j).append("\", ")
                            .append("\"startTime\": ").append(1725641066535L + j * 1000).append(", ")
                            .append("\"endTime\": ").append(1725641066535L + j * 1500).append(" }");

                    // 如果不是最后一个 span，添加逗号分隔符
                    if (j < 10) {
                        spansBuilder.append(", ");
                    }
                }
                spansBuilder.append("]"); // 结束 spans 数组

                // 构建完整的 JSON 数据
                String jsonData = "{ \"traceId\": \"" + traceId + "\", "
                        + "\"segmentId\": -1, "
                        + "\"serviceName\": \"service_" + i + "\", "
                        + "\"serviceId\": \"service_" + i + "\", "
                        + "\"spans\": " + spansBuilder.toString() + ", "  // 添加伪造的 spans 数据
                        + "\"startTime\": 1725641066535, "
                        + "\"endTime\": 1725641067000, "
                        + "\"peer\": \"peer_" + i + "\" }";

                // 将数据添加到批量请求中
                IndexRequest indexRequest = new IndexRequest("test")
                        .source(jsonData, XContentType.JSON);
                bulkRequest.add(indexRequest);

                // 每 batchSize 次提交一次
                if (i % batchSize == 0) {
                    BulkResponse bulkResponse = client.bulk(bulkRequest, RequestOptions.DEFAULT);
                    if (bulkResponse.hasFailures()) {
                        System.out.println("Bulk insert failed: " + bulkResponse.buildFailureMessage());
                    }
                    bulkRequest = new BulkRequest(); // 重置 BulkRequest 以处理下一批数据
                }
            }

            // 插入最后一批数据
            if (bulkRequest.numberOfActions() > 0) {
                BulkResponse bulkResponse = client.bulk(bulkRequest, RequestOptions.DEFAULT);
                if (bulkResponse.hasFailures()) {
                    System.out.println("Bulk insert failed: " + bulkResponse.buildFailureMessage());
                }
            }

            long end = System.currentTimeMillis();
            System.out.println("Inserted 9000000 records into Elasticsearch in " + (end - start) + " ms");
        }
    }
}
